from typing import Callable, ClassVar, Dict, Literal, cast, Tuple
from dataclasses import dataclass, field
import pandas as pd
import numpy as np
from . import base


@dataclass
class ChangeRate(base.BaseTransformer):
    new_col: str = 'pct'
    col: str = 'close'

    @classmethod
    def trans(cls, X: pd.Series):
        return X.pct_change()

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        data[self.new_col] = self.trans(data[self.col])
        return data


@dataclass
class OpenCloseChangRate(base.BaseTransformer):
    new_col: str = field(default='pct_oc')
    col_open: str = field(default='open')
    col_close: str = field(default='close')

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs):
        data = X.copy() if copy else X
        data[self.new_col] = np.where(
            data[self.col_close] - data[self.col_open] != 0,
            (data[self.col_close] - data[self.col_open]) / data[self.col_open],
            0
        )  # 如果是一字板，就认为是0
        return data


@dataclass
class IsIncreasement(base.BaseTransformer):
    new_col: str = field(default='is_inc')
    col: str = field(default='close')
    labels: Tuple[int, int, int] = (1, 0, -1)

    @classmethod
    def trans(cls, X: pd.Series, labels: Tuple[int, int, int]):
        i, p, d = labels
        return np.where(
            X > X.shift(1),
            i,
            np.where(X < X.shift(1), d, p)
        )

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        col = X[self.col]
        data[self.new_col] = self.trans(col, self.labels)
        return data


@dataclass
class ContinuousIncreasementCumsum(base.BaseTransformer):
    new_col: str = field(default='cont_inc_cumsum')
    col: str = field(default='close')
    col_inc: str = field(default='cont_inc')
    labels: Tuple[int, int, int] = (1, 0, -1)

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        if self.col_inc in data:
            inc = data[self.col_inc]
        else:
            inc = IsIncreasement.trans(data[self.col], self.labels)
        inc = cast(pd.Series, inc)
        data[self.new_col] = inc.cumsum()
        return data


@dataclass
class IsDecreasement(base.BaseTransformer):
    new_col: str = field(default='cont_dec')
    col: str = field(default='close')

    @classmethod
    def trans(cls, X: pd.Series):
        is_inc = X <= X.shift(1)
        return np.where(is_inc, 1, 0)

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        col = X[self.col]
        data[self.new_col] = self.trans(col)
        return data


@dataclass
class ContinuousDecreasementCumsum(base.BaseTransformer):
    new_col: str = field(default='cont_dec_cumsum')
    col: str = field(default='close')
    col_inc: str = field(default='cont_dec')

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        if self.col_inc in data:
            dec = data[self.col_inc]
        else:
            dec = IsDecreasement.trans(data[self.col])
        dec = cast(pd.Series, dec)
        data[self.new_col] = dec.cumsum()
        return data


@dataclass
class IsCapacityIncrease(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'cap_inc'
    col_inc_p: str = 'is_inc'
    col_inc_v: str = 'is_inc_vol'
    labels: Tuple[int, int] = (1, -1)
    inc_label: int = 1

    @classmethod
    def trans(cls, inc_p: pd.Series, inc_v: pd.Series, labels: Tuple[int, int], inc_label: int):
        i, d = labels
        return np.where(
            (inc_p == inc_label) & (inc_v == inc_label),
            i, d
        )

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col_inc_p], data[self.col_inc_v], self.labels, self.inc_label)
        return data


@dataclass
class IsCapacityDecrease(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'cap_dec'
    col_dec_p: str = 'is_dec'
    col_dec_v: str = 'is_dec_vol'
    labels: Tuple[int, int] = (1, -1)
    dec_label: int = -1

    @classmethod
    def trans(cls, dec_p: pd.Series, dec_v: pd.Series, labels: Tuple[int, int], dec_label: int):
        i, d = labels
        return np.where(
            (dec_p == dec_label) & (dec_v == dec_label),
            i, d
        )

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col_dec_p], data[self.col_dec_v], self.labels, self.dec_label)
        return data


@dataclass
class CumOp(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'cum'
    col: str = 'close'
    method: Literal['cumsum', 'cummax', 'cummin', 'cumprod'] = 'cumsum'

    METH_MAP: ClassVar[Dict[str, Callable[[pd.Series], pd.Series]]] = {
        'cumsum': lambda x: x.cumsum(),
        'cummax': lambda x: x.cummin(),
        'cummin': lambda x: x.cummin(),
        'cumprod': lambda x: x.cumprod(),
    }

    @classmethod
    def trans(cls, X: pd.Series, meth: str):
        return cls.METH_MAP[meth](X)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.METH_MAP[self.method](data[self.col])
        return data


@dataclass
class CumReturn(base.BaseTransformer):
    new_col: str = 'cumret'
    col: str = 'close'
    n: int = 10

    @classmethod
    def trans(cls, X: pd.Series, n: int):
        return X / X.shift(n) - 1

    def transform(self, X: pd.DataFrame, copy: bool = True, *args, **kwargs) -> pd.DataFrame:
        data = X.copy() if copy else X
        data[self.new_col] = self.trans(data[self.col], self.n)
        return data


@dataclass
class ForBackRatio(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'forbackratio'
    col: str = 'close'

    @classmethod
    def trans(cls, X: pd.Series):
        return X / X.shift(1)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = np.where(data[self.col].shift(1) > 0,
                                      self.trans(data[self.col]), 0)  # 当前一日是一字板时，就认为是0
        return data


@dataclass
class Bias(base.BaseTransformer, base.CopyMixin):
    new_col: str = 'bias'
    col: str = 'close'
    col_base: str = 'close'
    is_abs: bool = False

    @classmethod
    def trans(cls, target: pd.Series, x: pd.Series, is_abs: bool):
        return ((x - target) / target).abs() if is_abs else (x - target) / target

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(X[self.col_base], X[self.col], self.is_abs)
        return data


@dataclass
class RollingOp(base.BaseTransformer, base.CopyMixin):
    n: int = 20
    new_col: str = 'rollsum'
    col: str = 'close'
    method: Literal['sum', 'mean', 'max', 'min', 'std', 'first', 'last'] = 'sum'
    roll_args: dict = field(default_factory=dict)

    @classmethod
    def trans(cls, X: pd.Series, meth: str, *args, **kwargs):
        return X.rolling(*args, **kwargs).agg(meth)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], self.method, self.n, **self.roll_args)
        return data


@dataclass
class EmwOp(base.BaseTransformer, base.CopyMixin):
    span: int = 20
    new_col: str = 'emw'
    col: str = 'close'
    method: Literal['mean', 'std'] = 'mean'
    emw_args: dict = field(default_factory=dict)

    @classmethod
    def trans(cls, X: pd.Series, meth: str, *args, **kwargs):
        return X.ewm(*args, **kwargs).agg(meth)

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], span=self.span, meth=self.method, **self.emw_args)
        return data


@dataclass
class AnyChgPct(base.BaseTransformer, base.CopyMixin):
    n: int = 1
    new_col: str = 'pct'
    col: str = 'close'
    
    @classmethod
    def trans(cls, x: pd.Series, n: int):
        return x.shift(n) / x - 1 if n < 0 else x / x.shift(n) - 1
    
    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], self.n)
        return data


@dataclass
class AutoCorr(base.BaseTransformer, base.CopyMixin):
    new_col: str = "auto_corr"
    col: str = "close"
    n: int = 10
    lag: int = 1

    @classmethod
    def trans(cls, x: pd.Series, n: int, lag: int):
        return x.rolling(n).agg(lambda i: i.corr(i.shift(lag)))

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], self.n, self.lag)
        return data


@dataclass
class ForBackMeanComp(base.BaseTransformer, base.CopyMixin):
    new_col: str = "mean_comp"
    col: str = "close"
    n: int = 20
    head: int = 1

    @classmethod
    def trans(cls, x: pd.Series, n: int, head: int):
        return x.rolling(n).agg(lambda i: i.iloc[head:].mean() / (i.iloc[:head].mean() + 1/100000))

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = self.trans(data[self.col], self.n, self.head)
        return data


def getVol(close: pd.Series,
           value: int,
           unit: str = "minutes",
           span0: int = 100):
    """计算波动率
    :param close: 收盘价，index必须为时间戳
    """
    # daily vol, reindexed to close
    base_iloc = close.index.searchsorted(
        close.index-pd.Timedelta(value=value, unit=unit))
    base_iloc = base_iloc[base_iloc > 0]
    base_idx = pd.Series(close.index[base_iloc-1],
                         index=close.index[close.shape[0]-base_iloc.shape[0]:])
    df0 = close.loc[base_idx.index] / \
        close.loc[base_idx.values].values-1  # daily returns
    df0 = df0.ewm(span=span0).std()
    return df0
